import json
import time
from datetime import datetime
from colorama import Fore, Style, init
import copy

# 初始化colorama
init()

class LoggerTool:
    """日志工具类，用于打印调用agent信息和工具信息"""
    
    def __init__(self):
        self.start_time = time.time()
        self.last_time = self.start_time
        self.current_agent = None
        self.current_task = None
        # 添加工作流记录
        self.workflow_steps = []
        self.step_id = 0
        self.user_query = None
    
    def record_user_query(self, query):
        """记录用户查询"""
        self.user_query = query
        self.workflow_steps = []  # 重置工作流步骤
        self.step_id = 0
        self.start_time = time.time()
        self._add_workflow_step("user_input", None, "👤", query)
    
    def _add_workflow_step(self, type_name, agent_name, icon, content, operation=None, details=None, tool_name=None, inputs=None, outputs=None, from_agent=None, to_agent=None, task=None, reason=None):
        """添加工作流步骤"""
        self.step_id += 1
        timestamp = datetime.now().strftime("%Y-%m-%dT%H:%M:%S")
        elapsed = time.time() - self.start_time
        
        step = {
            "step_id": self.step_id,
            "timestamp": timestamp,
            "elapsed_time": round(elapsed, 2),
            "type": type_name,
            "icon": icon
        }
        
        if agent_name:
            step["agent"] = agent_name
        
        if operation:
            step["operation"] = operation
            
        if content:
            step["content"] = content
            
        if tool_name:
            step["tool_name"] = tool_name
            
        if inputs:
            step["input"] = inputs
            
        if outputs:
            step["output"] = outputs
            
        if details:
            step["details"] = details
            
        if from_agent:
            step["from_agent"] = from_agent
            
        if to_agent:
            step["to_agent"] = to_agent
            
        if task:
            step["task"] = task
            
        if reason:
            step["reason"] = reason
            
        self.workflow_steps.append(step)
        
    def get_workflow(self):
        """获取完整的工作流记录"""
        if not self.workflow_steps:
            return None
            
        # 获取所有参与的Agent
        agents_involved = set()
        tools_used = set()
        intent_type = None
        scene = None
        
        for step in self.workflow_steps:
            if step.get("agent"):
                agents_involved.add(step.get("agent"))
            if step.get("from_agent"):
                agents_involved.add(step.get("from_agent"))
            if step.get("to_agent"):
                agents_involved.add(step.get("to_agent"))
            if step.get("tool_name"):
                tools_used.add(step.get("tool_name"))
            if step.get("type") == "detection" and "意图类型" in step.get("content", ""):
                intent_type = step.get("content").split(":")[-1].strip()
            if step.get("type") == "detection" and "检测到场景" in step.get("content", ""):
                scene = step.get("content").split(":")[-1].strip()
                
        workflow = {
            "workflow_id": f"wf-{datetime.now().strftime('%Y%m%d-%H%M%S')}",
            "query": self.user_query,
            "timestamp": self.workflow_steps[0]["timestamp"] if self.workflow_steps else datetime.now().strftime("%Y-%m-%dT%H:%M:%S"),
            "total_execution_time": self.workflow_steps[-1]["elapsed_time"] if self.workflow_steps else 0,
            "status": "completed",
            "steps": self.workflow_steps,
            "workflow_summary": {
                "total_steps": len(self.workflow_steps),
                "agents_involved": list(agents_involved),
                "tools_used": list(tools_used)
            }
        }
        
        if intent_type:
            workflow["workflow_summary"]["detected_intent"] = intent_type
            
        if scene:
            workflow["workflow_summary"]["detected_scene"] = scene
            
        return workflow
    
    def log_agent_action(self, agent_name, action, details=None):
        """
        记录agent的行动
        
        参数:
            agent_name: agent的名称
            action: 执行的动作
            details: 动作的详细信息
        """
        self.current_agent = agent_name
        timestamp = datetime.now().strftime("%H:%M:%S")
        elapsed = time.time() - self.start_time
        
        print(f"\n{Fore.CYAN}[{timestamp}] ({elapsed:.2f}s) 🤖 Agent: {agent_name}{Style.RESET_ALL}")
        print(f"{Fore.BLUE}├─ 动作: {action}{Style.RESET_ALL}")
        
        details_copy = None
        if details:
            if isinstance(details, dict):
                details_str = json.dumps(details, ensure_ascii=False, indent=2)
                print(f"{Fore.BLUE}└─ 详情: \n{details_str}{Style.RESET_ALL}")
                details_copy = copy.deepcopy(details)
            else:
                print(f"{Fore.BLUE}└─ 详情: {details}{Style.RESET_ALL}")
                details_copy = details
                
        # 记录工作流步骤
        self._add_workflow_step("agent_action", agent_name, "🤖", None, None, details_copy, None, None, None, None, None, action)
    
    def log_workflow(self, message, level="INFO", agent_name=None):
        """
        记录工作流程信息
        
        参数:
            message: 要记录的消息
            level: 日志级别 (INFO, WARNING, ERROR)或操作类型
            agent_name: 执行该工作流的Agent名称
        """
        timestamp = datetime.now().strftime("%H:%M:%S")
        elapsed = time.time() - self.start_time
        
        # 根据level选择颜色
        color = Fore.GREEN
        if level == "WARNING":
            color = Fore.YELLOW
        elif level == "ERROR":
            color = Fore.RED
            
        # 根据不同情况选择适合的图标和组件名称
        if agent_name or self.current_agent:
            agent_info = f"[{agent_name or self.current_agent}] "
            # 根据agent角色选择适合的图标
            if "推荐" in (agent_name or self.current_agent):
                icon = "🛒"  # 推荐专家图标
            elif "库存" in (agent_name or self.current_agent):
                icon = "📦"  # 库存管理员图标
            elif "收银" in (agent_name or self.current_agent):
                icon = "💰"  # 收银员图标
            elif "总管理员" in (agent_name or self.current_agent):
                icon = "👨‍💼"  # 总管理员图标
            elif "分析" in (agent_name or self.current_agent) or "意图" in (agent_name or self.current_agent):
                icon = "🧠"  # 意图分析图标 
            else:
                icon = "🤖"  # 默认Agent图标
        else:
            # 根据消息内容选择合适的图标和组件名称
            if "分析用户意图" in message or "意图" in message:
                icon = "🧠"  # 意图分析图标
                agent_info = "[意图分析器] "
            elif "检测到场景" in message or "场景" in message:
                icon = "🔍"  # 场景检测图标
                agent_info = "[场景检测器] "
            elif "处理商品推荐" in message or "推荐" in message:
                icon = "🛒"  # 商品推荐图标
                agent_info = "[商品推荐专家] "
            elif "健身" in message or "运动" in message:
                icon = "💪"  # 健身场景图标
                agent_info = "[健身推荐专家] "
            elif "类别" in message or "分类" in message:
                icon = "🏷️"  # 类别图标
                agent_info = "[类别分析器] "
            elif "库存" in message or "商品查询" in message:
                icon = "📦"  # 库存查询图标
                agent_info = "[库存管理员] "
            elif "支付" in message or "付款" in message:
                icon = "💰"  # 支付处理图标
                agent_info = "[收银员] "
            elif "初始化" in message or "服务" in message or "启动" in message:
                icon = "🔧"  # 系统图标
                agent_info = "[系统] "
            elif "执行" in message or "任务" in message:
                icon = "⚙️"  # 任务执行图标
                agent_info = "[任务执行器] "
            elif "生成" in message or "文案" in message:
                icon = "✍️"  # 文案生成图标
                agent_info = "[文案生成器] "
            elif "CrewAI" in message or "Crew" in message:
                icon = "👥"  # Crew图标
                agent_info = "[Crew管理器] "
            elif "错误" in message or "出错" in message:
                icon = "⚠️"  # 错误图标
                agent_info = "[错误处理器] "
            else:
                # 即使是默认情况下也为消息分配一个明确的组件
                icon = "🔄"  # 默认处理图标
                agent_info = "[流程处理器] "
            
        # 将日志级别转换为更有描述性的操作名称
        action_type = level
        if level == "INFO":
            # 根据消息内容判断操作类型
            if "分析" in message:
                action_type = "分析"
            elif "检测" in message:
                action_type = "检测"
            elif "处理" in message:
                action_type = "处理"
            elif "使用" in message:
                action_type = "使用"
            elif "生成" in message:
                action_type = "生成"
            elif "初始化" in message or "启动" in message:
                action_type = "初始化"
            elif "执行" in message:
                action_type = "执行"
            elif "完成" in message:
                action_type = "完成"
            else:
                action_type = "操作"
        elif level == "WARNING":
            action_type = "警告"
        elif level == "ERROR":
            action_type = "错误"
            
        task_info = f"({self.current_task}) " if self.current_task else ""
        print(f"\n{color}[{timestamp}] ({elapsed:.2f}s) {icon} {action_type}: {agent_info}{task_info}{message}{Style.RESET_ALL}")
        
        # 确定用于工作流的类型
        workflow_type = "processing"  # 默认类型
        if "分析" in action_type:
            workflow_type = "analysis"
        elif "检测" in action_type:
            workflow_type = "detection"
        elif "错误" in action_type:
            workflow_type = "error"
        
        # 记录工作流步骤
        real_agent = None
        if agent_name:
            real_agent = agent_name
        elif self.current_agent:
            real_agent = self.current_agent
        elif agent_info and agent_info.startswith("[") and agent_info.endswith("] "):
            real_agent = agent_info[1:-2]  # 提取中括号内的Agent名称
            
        self._add_workflow_step(workflow_type, real_agent, icon, message, action_type)
    
    def log_tool_usage(self, tool_name, inputs=None, outputs=None, agent_name=None):
        """
        记录工具的使用情况
        
        参数:
            tool_name: 工具的名称
            inputs: 工具的输入参数
            outputs: 工具的输出结果
            agent_name: 使用该工具的Agent名称
        """
        timestamp = datetime.now().strftime("%H:%M:%S")
        elapsed = time.time() - self.start_time
        
        # 如果没有提供agent_name，根据工具名称推断
        if not agent_name and not self.current_agent:
            if "意图" in tool_name or "nlp" in tool_name.lower():
                agent_name = "意图分析器"
            elif "推荐" in tool_name or "健身" in tool_name:
                agent_name = "商品推荐专家"
            elif "库存" in tool_name or "商品查询" in tool_name:
                agent_name = "库存管理员"
            elif "支付" in tool_name or "付款" in tool_name:
                agent_name = "收银员"
            elif "场景" in tool_name:
                agent_name = "场景检测器"
            elif "类别" in tool_name or "分类" in tool_name:
                agent_name = "类别分析器"
            elif "价格" in tool_name:
                agent_name = "价格分析器"
            elif "搜索" in tool_name or "查询" in tool_name:
                agent_name = "搜索引擎"
            else:
                agent_name = "工具调用者"  # 默认名称
        
        agent_info = f"[{agent_name or self.current_agent}] " if (agent_name or self.current_agent) else ""
        
        # 根据工具名称选择适当的图标
        if "意图" in tool_name:
            icon = "🧠"  # 意图分析图标
        elif "推荐" in tool_name or "健身" in tool_name:
            icon = "🛒"  # 推荐图标
        elif "库存" in tool_name:
            icon = "📦"  # 库存图标
        elif "支付" in tool_name:
            icon = "💰"  # 支付图标
        elif "场景" in tool_name:
            icon = "🔍"  # 场景图标
        elif "搜索" in tool_name or "查询" in tool_name:
            icon = "🔎"  # 搜索图标
        elif "价格" in tool_name:
            icon = "💲"  # 价格图标
        else:
            icon = "🔧"  # 默认工具图标
            
        print(f"\n{Fore.YELLOW}[{timestamp}] ({elapsed:.2f}s) {icon} 工具调用: {agent_info}{tool_name}{Style.RESET_ALL}")
        
        inputs_copy = None
        if inputs:
            if isinstance(inputs, dict):
                inputs_str = json.dumps(inputs, ensure_ascii=False, indent=2)
                print(f"{Fore.YELLOW}├─ 输入: \n{inputs_str}{Style.RESET_ALL}")
                inputs_copy = copy.deepcopy(inputs)
            else:
                print(f"{Fore.YELLOW}├─ 输入: {inputs}{Style.RESET_ALL}")
                inputs_copy = inputs
        
        outputs_copy = None
        if outputs:
            if isinstance(outputs, dict):
                outputs_str = json.dumps(outputs, ensure_ascii=False, indent=2)
                print(f"{Fore.YELLOW}└─ 输出: \n{outputs_str}{Style.RESET_ALL}")
                outputs_copy = copy.deepcopy(outputs)
            else:
                print(f"{Fore.YELLOW}└─ 输出: {outputs}{Style.RESET_ALL}")
                outputs_copy = outputs
                
        # 记录工作流步骤
        real_agent = agent_name or self.current_agent
        if agent_info and not real_agent:
            if agent_info.startswith("[") and agent_info.endswith("] "):
                real_agent = agent_info[1:-2]  # 提取中括号内的Agent名称
                
        self._add_workflow_step("tool_call", real_agent, icon, None, "工具调用", None, tool_name, inputs_copy, outputs_copy)
    
    def log_user_interaction(self, message):
        """
        记录用户交互信息
        
        参数:
            message: 用户消息
        """
        timestamp = datetime.now().strftime("%H:%M:%S")
        print(f"\n{Fore.MAGENTA}[{timestamp}] 👤 用户: {message}{Style.RESET_ALL}")
        
        # 记录用户查询
        self.record_user_query(message)
    
    def log_system_response(self, response, agent_name=None):
        """
        记录系统响应信息
        
        参数:
            response: 系统响应内容
            agent_name: 生成响应的Agent名称
        """
        timestamp = datetime.now().strftime("%H:%M:%S")
        elapsed = time.time() - self.start_time
        
        # 如果没有提供agent_name，尝试从上下文推断
        if not agent_name:
            agent_name = self.current_agent
            
            # 如果没有current_agent，尝试从响应内容推断
            if not agent_name and isinstance(response, str):
                if "健身" in response or "运动" in response:
                    agent_name = "健身推荐专家"
                elif "商品" in response or "推荐" in response:
                    agent_name = "商品推荐专家"
                elif "库存" in response or "缺货" in response:
                    agent_name = "库存管理员"
                elif "支付" in response or "付款" in response:
                    agent_name = "收银员"
                elif "您好" in response or "欢迎" in response:
                    agent_name = "超市总管理员"
                else:
                    agent_name = "系统响应"
                    
        agent_info = f"[{agent_name}]" if agent_name else ""
        
        # 根据agent选择合适的图标
        if agent_name:
            if "健身" in agent_name:
                icon = "💪"  # 健身专家图标
            elif "推荐" in agent_name:
                icon = "🛒"  # 推荐专家图标
            elif "库存" in agent_name:
                icon = "📦"  # 库存管理员图标
            elif "收银" in agent_name:
                icon = "💰"  # 收银员图标
            elif "总管理员" in agent_name:
                icon = "👨‍💼"  # 总管理员图标
            else:
                icon = "🤖"  # 默认系统图标
        else:
            icon = "🤖"  # 默认系统图标
            
        print(f"\n{Fore.GREEN}[{timestamp}] ({elapsed:.2f}s) {icon} 系统响应: {agent_info}{Style.RESET_ALL}")
        print(f"{Fore.GREEN}{response}{Style.RESET_ALL}")
        
        # 记录工作流步骤
        response_content = response
        if isinstance(response, dict):
            response_content = json.dumps(response, ensure_ascii=False)
            
        self._add_workflow_step("system_response", agent_name, icon, response_content)

    def log_agent_delegation(self, from_agent, to_agent, task, reason=None):
        """
        记录Agent任务委派
        
        参数:
            from_agent: 委派任务的Agent
            to_agent: 接收任务的Agent
            task: 委派的任务内容
            reason: 委派原因(可选)
        """
        timestamp = datetime.now().strftime("%H:%M:%S")
        elapsed = time.time() - self.start_time
        print(f"\n{Fore.CYAN}[{timestamp}] ({elapsed:.2f}s) 📤 任务委派: {from_agent} -> {to_agent}{Style.RESET_ALL}")
        print(f"{Fore.BLUE}├─ 任务: {task}{Style.RESET_ALL}")
        if reason:
            print(f"{Fore.BLUE}└─ 原因: {reason}{Style.RESET_ALL}")
            
        # 记录工作流步骤
        self._add_workflow_step("task_delegation", None, "📤", None, None, None, None, None, None, from_agent, to_agent, task, reason)

    def log_agent_thinking(self, agent_name, thought):
        """
        记录agent的思考过程
        
        参数:
            agent_name: agent的名称
            thought: 思考的内容
        """
        timestamp = datetime.now().strftime("%H:%M:%S")
        elapsed = time.time() - self.start_time
        
        print(f"\n{Fore.MAGENTA}[{timestamp}] ({elapsed:.2f}s) 🤔 思考: [{agent_name}]{Style.RESET_ALL}")
        print(f"{Fore.MAGENTA}└─ {thought}{Style.RESET_ALL}")
        
        # 记录工作流步骤
        self._add_workflow_step("thinking", agent_name, "🤔", thought)